#include "neuralquantizer.h"

namespace WingGif {

NeuralQuantizer::NeuralQuantizer(int samplingFactor, int maximumColors)
    : Quantizer(false, this) {
  _samplingFactor = samplingFactor;
  _networkSize = maximumColors;
}

void NeuralQuantizer::firstPass(QByteArray pixels) {
  auto maxColorsWithTransparency =
      transparentColor().isValid() ? _networkSize - 1 : _networkSize;
  setMaxColorsWithTransparency(maxColorsWithTransparency);

  _maximumNeuronIndex = maxColorsWithTransparency - 1;
  _network.resize(maxColorsWithTransparency);
  _indexOfGreen.resize(256);
  _biases.resize(maxColorsWithTransparency);
  _frequencies.resize(maxColorsWithTransparency);
  _initialNeighbourhoodSize = qMax(maxColorsWithTransparency >> 3, 1);
  _neighbourhoodAlphas.resize(_initialNeighbourhoodSize);
  _initialUnbiasedNeighbourhoodSize =
      _initialNeighbourhoodSize * neighbourhoodSizeBias;

  for (auto neuronIndex = 0; neuronIndex < maxColorsWithTransparency;
       neuronIndex++) {
    _network[neuronIndex][0] = _network[neuronIndex][1] =
        _network[neuronIndex][2] =
            (neuronIndex << (networkBiasShift + 8)) / maxColorsWithTransparency;

    _frequencies[neuronIndex] = intBias / maxColorsWithTransparency;
    _biases[neuronIndex] = 0;
  }

  learn(pixels);
  unbiasNetwork();
  buildIndex();
}

QVector<QColor> NeuralQuantizer::buildPalette() {
  auto _maxColorsWithTransparency = maxColorsWithTransparency();

  auto map = new char[ulong(3 * _maxColorsWithTransparency)];
  auto index = new int[ulong(_maxColorsWithTransparency)];

  // Gets the index of each color.
  for (auto i = 0; i < _maxColorsWithTransparency; i++)
    index[_network[i][3]] = i;

  QVector<QColor> colors;

  auto k = 0;
  for (auto i = 0; i < _maxColorsWithTransparency; i++) {
    auto j = index[i];

    // BGR.
    map[k++] = char(_network[j][0]);
    map[k++] = char(_network[j][1]);
    map[k++] = char(_network[j][2]);

    // Add repeated colors?
    colors.append(
        QColor::fromRgba(qRgba(map[k - 1], map[k - 2], map[k - 3], 255)));
  }

  auto _transparentColor = transparentColor();
  if (_transparentColor.isValid())
    colors.append(_transparentColor);

  delete[] map;
  delete[] index;
  return colors;
}

uchar NeuralQuantizer::quantizePixel(QColor pixel) {
  return mapColor(pixel.blue(), pixel.green(), pixel.red());
}

void NeuralQuantizer::learn(QByteArray pixels) {
  _pixelBytesCount = pixels.length();

  // Preparations for learning

  // If the image is so small that it has fewer pixels than the largest prime
  // number used to determine how to step through the pixels, include every
  // pixel in the sample.
  if (_pixelBytesCount < prime4 * 4)
    _samplingFactor = 1;

  auto alphaDecrement = 30 + (_samplingFactor - 1) / 4;
  auto pixelIndex = 0;

  // Set the number of elements of the learning data to be examined during the
  // learning loop. Pixels are in BGRA. If _samplingFactor is 1 then every
  // element will be examined. If _samplingFactor is 10 then one tenth of the
  // elements will be examined.
  auto pixelsToExamine = _pixelBytesCount / (4 * _samplingFactor);

  // Set how often the alpha value for shifting neurons is updated.
  // A value of 1 means it is updated once per pixel examined, 10 means it is
  // updated every 10 pixels, and so on.
  auto alphaUpdateFrequency = qMax(1, pixelsToExamine / numberOfLearningCycles);

  // Alpha is a factor which controls how far neurons are moved during the
  // learning loop, and it decreases as learning proceeds.
  auto alpha = initialAlpha;

  // Set the size of the neighbourhood which makes up the neighbouring neurons
  // which also need to be moved when a neuron is moved.
  auto unbiasedNeighbourhoodSize = _initialUnbiasedNeighbourhoodSize;

  auto neighbourhoodSize =
      unbiasedNeighbourhoodSize >> neighbourhoodSizeBiasShift;

  // Is this possible?
  if (neighbourhoodSize < 1)
    neighbourhoodSize = 1;

  // Set the initial alpha values for neighbouring neurons.
  setNeighbourhoodAlphas(_neighbourhoodAlphas, neighbourhoodSize, alpha,
                         radiusBias);

  // Get the number of pixels to skip between samples.
  auto step = getPixelIndexIncrement(_pixelBytesCount);

  // Learning

  auto pixelsExamined = 0;

  // var hashTable = new HashSet<int>();

  while (pixelsExamined < pixelsToExamine) {
    // By trying to ignore repeated colors, this gives the opportunity to other
    // colors to be used instead, which may cause visual imperfections,
    // specially with full frames (with lots of colors). var hash =
    // BitConverter.ToInt32(new[] { byte.MaxValue, pixels[pixelIndex + 0],
    // pixels[pixelIndex + 1], pixels[pixelIndex + 2] }, 0);

    // Only ignore transparent pixels.
    if (pixels[pixelIndex + 3] > char(0)) // && !hashTable.Contains(hash))
    {
      // Move neurons

      auto blue = (pixels[pixelIndex + 0] & 0xff) << networkBiasShift;
      auto green = (pixels[pixelIndex + 1] & 0xff) << networkBiasShift;
      auto red = (pixels[pixelIndex + 2] & 0xff) << networkBiasShift;

      auto bestNeuronIndex = findClosestAndReturnBestNeuron(blue, green, red);

      // Move this neuron closer to the current element of the learning data by
      // a factor of alpha.
      moveNeuron(alpha, bestNeuronIndex, blue, green, red);

      // If appropriate, move neighbouring neurons closer to the color of the
      // current pixel.
      if (neighbourhoodSize != 0)
        moveNeighbouringNeurons(neighbourhoodSize, bestNeuronIndex, blue, green,
                                red);
    }

    // hashTable.Add(hash);

    // Move on to the next learning data element to be examined

    pixelIndex += step;

    // If gone past the end of the learning data, wrap around to the start
    // again.
    if (pixelIndex >= _pixelBytesCount)
      pixelIndex -= _pixelBytesCount;

    // Keep track of how many elements have been examined so far.
    pixelsExamined++;

    // Update the alpha values for moving neurons if appropriate

    if (pixelsExamined % alphaUpdateFrequency == 0) {
      alpha -= alpha / alphaDecrement;
      unbiasedNeighbourhoodSize -=
          unbiasedNeighbourhoodSize / unbiasedNeighbourhoodSizeDecrement;
      neighbourhoodSize =
          unbiasedNeighbourhoodSize >> neighbourhoodSizeBiasShift;

      if (neighbourhoodSize <= 1)
        neighbourhoodSize = 0;

      // Update the alpha values to be used for moving neighbouring neurons.
      setNeighbourhoodAlphas(_neighbourhoodAlphas, neighbourhoodSize, alpha,
                             radiusBias);
    }
  }
}

void NeuralQuantizer::setNeighbourhoodAlphas(QVector<int> neighbourhoodAlphas,
                                             int neighbourhoodSize, int alpha,
                                             int radiusBias) {
  // Get neighbourhood size squared - only need to calculate this once.
  auto squared = neighbourhoodSize * neighbourhoodSize;

  for (auto i = 0; i < neighbourhoodSize; i++)
    neighbourhoodAlphas[i] = alpha * ((squared - i * i) * radiusBias / squared);
}

int NeuralQuantizer::getPixelIndexIncrement(int pictureByteCount) {
  int step;

  if (pictureByteCount < prime4 * 4)
    step = 4;
  else if (pictureByteCount % prime1 !=
           0) // The number of pixels is not divisible by the first prime
              // number.
    step = prime1 * 4;
  else if (pictureByteCount % prime2 !=
           0) // The number of pixels is not divisible by the second prime
              // number.
    step = prime2 * 4;
  else if (pictureByteCount % prime3 !=
           0) // The number of pixels is not divisible by the third prime
              // number.
    step = prime3 * 4;
  else {
    // The number of pixels is divisible by the first, second and third prime
    // numbers. To cover this in a test case we'd need learning data consisting
    // of over 119 million neurons!
    step = prime4 * 4;
  }

  return step;
}

int NeuralQuantizer::findClosestAndReturnBestNeuron(int blue, int green,
                                                    int red) {
  auto bestDistance = ~(1u << 31); // Bitwise inverted.
  auto bestBiasDistance = bestDistance;
  auto closestNeuronIndex = -1;
  auto bestBiasNeuronIndex = closestNeuronIndex;
  auto _maxColorsWithTransparency = maxColorsWithTransparency();

  for (auto neuronIndex = 0; neuronIndex < _maxColorsWithTransparency;
       neuronIndex++) {
    // Calculate the distance

    auto distance = _network[neuronIndex][0] - blue;

    if (distance < 0)
      distance = -distance;

    auto distanceIncrement = _network[neuronIndex][1] - green;

    if (distanceIncrement < 0)
      distanceIncrement = -distanceIncrement;

    distance += distanceIncrement;
    distanceIncrement = _network[neuronIndex][2] - red;

    if (distanceIncrement < 0)
      distanceIncrement = -distanceIncrement;

    distance += distanceIncrement;

    if (uint(distance) < bestDistance) {
      bestDistance = uint(distance);
      closestNeuronIndex = neuronIndex;
    }

    // Calculate the bias distance

    // Bias distance takes into account the distance between the neuron and the
    // coordinate, and also the neuron's bias. The more frequently a neuron has
    // already been chosen, the lower its bias, so less frequently-chosen
    // neurons have a better chance of being returned by this method. This
    // ensures that the distribution of neurons is densest in areas of the
    // network space which have most coordinates in the learning data.
    auto biasDistance =
        distance - (_biases[neuronIndex] >> (intBiasShift - networkBiasShift));

    if (uint(biasDistance) < bestBiasDistance) {
      bestBiasDistance = uint(biasDistance);
      bestBiasNeuronIndex = neuronIndex;
    }

    // Decrease the frequency and increase the bias for all neurons

    auto betaFrequency = _frequencies[neuronIndex] >> betaShift;

    _frequencies[neuronIndex] -= betaFrequency;
    _biases[neuronIndex] += betaFrequency << gammaShift;
  }

  // Increase the frequency and decrease the bias for just the closest neuron.
  _frequencies[closestNeuronIndex] += closestNeuronFrequencyIncrement;
  _biases[closestNeuronIndex] -= closestNeuronBiasDecrement;

  return bestBiasNeuronIndex;
}

void NeuralQuantizer::moveNeuron(int alpha, int neuronIndexToMove, int blue,
                                 int green, int red) {
  _network[neuronIndexToMove][0] -=
      (alpha * (_network[neuronIndexToMove][0] - blue)) / initialAlpha;
  _network[neuronIndexToMove][1] -=
      (alpha * (_network[neuronIndexToMove][1] - green)) / initialAlpha;
  _network[neuronIndexToMove][2] -=
      (alpha * (_network[neuronIndexToMove][2] - red)) / initialAlpha;
}

void NeuralQuantizer::moveNeighbouringNeurons(int neighbourhoodSize,
                                              int neuronIndex, int blue,
                                              int green, int red) {
  // Set lower and upper bounds of the neighbourhood of neurons to be moved

  auto lowNeuronIndexLimit = neuronIndex - neighbourhoodSize;

  if (lowNeuronIndexLimit < -1)
    lowNeuronIndexLimit = -1;

  auto highNeuronIndexLimit = neuronIndex + neighbourhoodSize;

  if (highNeuronIndexLimit > _network.length())
    highNeuronIndexLimit = _network.length();

  // Start with the neurons immediately before and after the specified index and
  // work outwards.
  auto highNeuronIndex = neuronIndex + 1;
  auto lowNeuronIndex = neuronIndex - 1;
  auto neighbourAlphaIndex = 1;

  while (highNeuronIndex < highNeuronIndexLimit ||
         lowNeuronIndex > lowNeuronIndexLimit) {
    auto neighbourhoodAlpha = _neighbourhoodAlphas[neighbourAlphaIndex++];

    if (highNeuronIndex < highNeuronIndexLimit)
      moveNeighbour(highNeuronIndex++, neighbourhoodAlpha, alphaRadiusBias,
                    blue, green, red);

    if (lowNeuronIndex > lowNeuronIndexLimit)
      moveNeighbour(lowNeuronIndex--, neighbourhoodAlpha, alphaRadiusBias, blue,
                    green, red);
  }
}

void NeuralQuantizer::moveNeighbour(int neuronIndexToMove, int alpha,
                                    int alphaRadiusBias, int blue, int green,
                                    int red) {
  _network[neuronIndexToMove][0] -=
      (alpha * (_network[neuronIndexToMove][0] - blue)) / alphaRadiusBias;
  _network[neuronIndexToMove][1] -=
      (alpha * (_network[neuronIndexToMove][1] - green)) / alphaRadiusBias;
  _network[neuronIndexToMove][2] -=
      (alpha * (_network[neuronIndexToMove][2] - red)) / alphaRadiusBias;
}

void NeuralQuantizer::unbiasNetwork() {
  for (auto neuronIndex = 0; neuronIndex < maxColorsWithTransparency();
       neuronIndex++) {
    _network[neuronIndex][0] >>= networkBiasShift;
    _network[neuronIndex][1] >>= networkBiasShift;
    _network[neuronIndex][2] >>= networkBiasShift;
    _network[neuronIndex][3] = neuronIndex; // Record the color number.
  }
}

void NeuralQuantizer::buildIndex() {
  int greenValue;
  auto previousLeastGreenValue = 0;
  auto startingGreenValue = 0;

  for (auto thisNeuronIndex = 0; thisNeuronIndex < maxColorsWithTransparency();
       thisNeuronIndex++) {
    auto thisNeuron = _network[thisNeuronIndex];

    // Find the least green neuron between the current neuron and the end of the
    // network.
    auto _indexOfLeastGreenNeuron = indexOfLeastGreenNeuron(thisNeuronIndex);
    auto leastGreenNeuron = _network[_indexOfLeastGreenNeuron];
    auto greenValueOfLeastGreenNeuron = leastGreenNeuron[1];

    // Move the neuron with the lowest index towards the beginning of the array.
    if (thisNeuronIndex != _indexOfLeastGreenNeuron)
      thisNeuron.swap(leastGreenNeuron);

    if (greenValueOfLeastGreenNeuron != previousLeastGreenValue) {
      // Then we've found a new least green neuron so update the array of green
      // indices accordingly
      _indexOfGreen[previousLeastGreenValue] =
          (startingGreenValue + thisNeuronIndex) >> 1;

      for (greenValue = previousLeastGreenValue + 1;
           greenValue < greenValueOfLeastGreenNeuron; greenValue++)
        _indexOfGreen[greenValue] = thisNeuronIndex;

      previousLeastGreenValue = greenValueOfLeastGreenNeuron;
      startingGreenValue = thisNeuronIndex;
    }
  }

  _indexOfGreen[previousLeastGreenValue] =
      (startingGreenValue + _maximumNeuronIndex) >> 1;

  // Fill the remainder of the _indexOfGreen array with the index of the last
  // neuron in the network.
  for (greenValue = previousLeastGreenValue + 1; greenValue < 256; greenValue++)
    _indexOfGreen[greenValue] = _maximumNeuronIndex;
}

int NeuralQuantizer::indexOfLeastGreenNeuron(int startNeuronIndex) {
  // Start with the current neuron, its index and green value.
  auto indexOfLeastGreenNeuron = startNeuronIndex;
  auto greenValueOfLeastGreenNeuron = _network[startNeuronIndex][1];

  // And compare it with the remaining neurons.
  for (auto otherNeuronIndex = startNeuronIndex + 1;
       otherNeuronIndex < maxColorsWithTransparency(); otherNeuronIndex++) {
    auto otherNeuron = _network[otherNeuronIndex];

    if (otherNeuron[1] < greenValueOfLeastGreenNeuron) {
      // The green value of otherNeuron is lower than that of the least green
      // neuron seen so far, so otherNeuron becomes the least green.
      indexOfLeastGreenNeuron = otherNeuronIndex;
      greenValueOfLeastGreenNeuron = otherNeuron[1];
    }
  }

  return indexOfLeastGreenNeuron;
}

uchar NeuralQuantizer::mapColor(int blue, int green, int red) {
  auto bestIndex = -1;
  auto bestDistance = 1000; // Biggest possible dist is 256 * 3.
  auto highNeuronIndex = _indexOfGreen[green]; // Index on g.
  auto lowNeuronIndex =
      highNeuronIndex - 1; // Start at netindex[g] and work outwards.
  auto _maxColorsWithTransparency = maxColorsWithTransparency();

  while (highNeuronIndex < _maxColorsWithTransparency || lowNeuronIndex >= 0) {
    int distance;
    QVector<int> thisNeuron;
    int distanceIncrement;

    if (highNeuronIndex < _maxColorsWithTransparency) {
      thisNeuron = _network[highNeuronIndex];
      thisNeuron.resize(4);
      distance = thisNeuron[1] - green; // Index key.

      if (distance >= bestDistance) {
        highNeuronIndex = _maxColorsWithTransparency; // Stop iteration.
      } else {
        highNeuronIndex++;

        if (distance < 0)
          distance = -distance;

        distanceIncrement = thisNeuron[0] - blue;

        if (distanceIncrement < 0)
          distanceIncrement = -distanceIncrement;

        distance += distanceIncrement;

        if (distance < bestDistance) {
          distanceIncrement = thisNeuron[2] - red;

          if (distanceIncrement < 0)
            distanceIncrement = -distanceIncrement;

          distance += distanceIncrement;

          if (distance < bestDistance) {
            bestDistance = distance;
            bestIndex = thisNeuron[3];
          }
        }
      }
    }

    if (lowNeuronIndex >= 0) {
      thisNeuron = _network[lowNeuronIndex];
      thisNeuron.resize(4);
      distance = green - thisNeuron[1]; // Index key.

      if (distance >= bestDistance) {
        lowNeuronIndex = -1; // Stop iteration.
      } else {
        lowNeuronIndex--;

        if (distance < 0)
          distance = -distance;

        distanceIncrement = thisNeuron[0] - blue;

        if (distanceIncrement < 0)
          distanceIncrement = -distanceIncrement;

        distance += distanceIncrement;

        if (distance < bestDistance) {
          distanceIncrement = thisNeuron[2] - red;

          if (distanceIncrement < 0)
            distanceIncrement = -distanceIncrement;

          distance += distanceIncrement;

          if (distance < bestDistance) {
            bestDistance = distance;
            bestIndex = thisNeuron[3];
          }
        }
      }
    }
  }

  return uchar(qMin(bestIndex, _maxColorsWithTransparency));
}

} // namespace WingGif
